import os.path

from transformers import TrainingArguments, Trainer  # 导入trainer相关的包
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from datasets import load_dataset
from dual_model import DualModel
import evaluate


class DualModelTrainer(object):
    #  "/data/datasets/SimCLUE/datasets/train_pair_1w.json" ,
    #  '/data/models/huggingface/chinese-macbert-base'
    def __init__(self, dataset_dir, pre_train_model_path,model_save_dir,train_model_dir=None):
        self.dataset_dir = dataset_dir
        self.pre_train_model_path = pre_train_model_path
        self.model_save_dir = model_save_dir
        """
        加载分词器
        """
        self.tokenizer = AutoTokenizer.from_pretrained(self.pre_train_model_path)
        # 数据集处理
        self.data_handler()
        """
        加载预训练模型
        """
        self.model  =  self.load_pretrain_model(train_model_dir)
        self.model = DualModel.from_pretrained(self.pre_train_model_path)

        """
        加载评估函数
        """
        self.acc_metric = evaluate.load('/evaluate-main/metrics/accuracy')
        self.f1_metric = evaluate.load('/evaluate-main/metrics/f1')

    def load_pretrain_model(self, train_model_dir):
        if train_model_dir:
            if os.path.exists(train_model_dir):
                print("继续训练")
                self.model = DualModel.from_pretrained(train_model_dir)
            else:
                self.model = DualModel.from_pretrained(self.pre_train_model_path)
        else:
            self.model = DualModel.from_pretrained(self.pre_train_model_path)

    def data_handler(self):
        """
        数据集加载，该数据集是一个1w数据的数据集，如果想提高匹配的准确度，最好更换数据集或者采用自定义数据集
        可以提高匹配准确率
        """
        # dataset_dir = "/data/datasets/SimCLUE/datasets/train_pair_1w.json"
        dataset = load_dataset('json', data_files=self.dataset_dir, split='train')
        """
        切分数据集合
        """
        self.datasets = dataset.train_test_split(test_size=0.05,shuffle=False)
        """
        数据处理
        """
        self.tokenizer_datasets = self.datasets.map(self.process_function, batched=True,
                                                    remove_columns=self.datasets['train'].column_names)

    """
    数据处理
    """

    def process_function(self, examples):
        sententces = []
        labels = []
        """
        将小批量数据集 examples 的 sentence1  sentence2 label 取出来。
        将文本加入到  sententces 列表中，将标注加到 labels 中
        """
        for sen1, sen2, label in zip(examples['sentence1'], examples['sentence2'], examples['label']):
            sententces.append(sen1)
            sententces.append(sen2)
            labels.append(1 if int(label) == 1 else -1)
        """
        由于每一个批次是2000条数据，将这些数据转化为词向量将将产生 key 为 ['input_ids', 'token_type_ids', 'attention_mask']的三个列表
        每一个列表信息都是2000条数据
        """
        tokenizer_examples = self.tokenizer(sententces, max_length=128, truncation=True, padding="max_length")
        """
        将2000条数据中每一个一维列表 都变成 一个二维列表，压缩为1000条数据
        input_ids 的 2000条，变为 input_ids 二维列表1000条
        token_type_ids 的 2000条，变为 input_ids 二维列表1000条
        attention_mask 的 2000条，变为 input_ids 二维列表1000条
        """
        tokenizer_examples = {k: [v[i:i + 2] for i in range(0, len(v), 2)] for k, v in tokenizer_examples.items()}
        # print(tokenizer_examples["attention_mask"][0])
        #  [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]
        # print(len(tokenizer_examples["input_ids"])) 1000
        """
        将labels 1000 条数据赋值到一个labels 为key的字典中去
        """
        tokenizer_examples["labels"] = labels
        return tokenizer_examples

    """
    写一个多种标准合成的评价函数
    """

    def eval_metrics(self, eval_predict):
        predictions, labels = eval_predict
        predictions = [int(p > 0.6) for p in predictions]
        # l 大于 0 为 1 小于 0 为 0 由于原始标签是 1和0 ,但是计算CosineSimilarity之后的标签是 -1 1 所以需要转换
        labels = [int(l > 0) for l in labels]

        combine_acc = self.acc_metric.compute(predictions=predictions, references=labels)
        f1 = self.f1_metric.compute(predictions=predictions, references=labels)
        combine_acc.update(f1)
        return combine_acc

    def train(self):
        """
        加载训练参数
        """
        train_args = TrainingArguments(
            output_dir=self.model_save_dir,  # 训练模型的输出路径
            per_device_train_batch_size=8,  # 训练集 batch size
            per_device_eval_batch_size=8,  # 验证集 batch size
            num_train_epochs=12,  # 训练的 epochs 数
            logging_steps=100,  # 每 150 个 batch 记录一次日志
            save_steps=100, # 每 100 个 batch 保存一次模型 和每个epoch 保存一次模型互斥
            save_strategy='steps',  # 每 1 个 epoch 保存一次模型
            # save_total_limit=20,  # 保存 1 个模型
            evaluation_strategy='epoch',  # 每 1 个 epoch 验证一次模型
            learning_rate=1e-5,  # 学习率
            weight_decay=0.01,  # 权重衰减
            metric_for_best_model='f1',
            # load_best_model_at_end=True,  # 加载最优模型
            # report_to='tensorboard',  # 训练日志保存到 tensorboard
        )
        """
        创建训练器
        """
        trainer = Trainer(
            model=self.model,  # 指定模型
            args=train_args,  # 执行训练参数
            train_dataset=self.tokenizer_datasets['train'],  # 训练数据集指定 切分后的数据集
            eval_dataset=self.tokenizer_datasets['test'],  # 验证数据集指定
            compute_metrics=self.eval_metrics,  # tokenizer
        )
        """
        开始训练
        """
        trainer.train()


def main():
    """
    1、训练数据优化。
    2、预训练模型优化
    3、训练参数调整（学习率）
    """
    dataset_dir = "/data/datasets/SimCLUE/datasets/train_pair_4w.json"
    # dataset_dir = "/data/datasets/SimCLUE/simclue_public/train_pair.json"
    pre_train_model_path = '/data/models/huggingface/chinese-macbert-large'
    # train_model_dir = "D:/code/logs/dual_model/checkpoints/checkpoint-1250"
    model_save_dir = '/data/logs/qq_search_robot/dual_model'
    DualModelTrainer(dataset_dir, pre_train_model_path,model_save_dir,None).train()

if __name__ == "__main__":
    main()
